import traceback
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse

# ---- LangChain 相关 ----
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder  # 用于消息占位的 ChatPrompt
from langchain.memory import ConversationBufferMemory                # 基于“缓冲区”的会话记忆
from langchain_ollama import OllamaLLM

# =========================================
# 这段代码里 ConversationBufferMemory 是如何“保留历史”的？
# -----------------------------------------
# 1) memory = ConversationBufferMemory(...):
#    - 内部维护一个“对话缓冲（buffer）”，以消息列表的形式存放 Human/AI 的来回对话。
#    - 我们设置 return_messages=True，使得 memory.load_memory_variables() 返回的是
#      [HumanMessage(...), AIMessage(...), ...] 的消息列表，而不是拼接后的纯字符串。
#
# 2) LLMChain 调用 predict() 时会做两件关键事：
#    a. 先调用 memory.load_memory_variables({当前输入})，拿到历史消息：
#         {"history": [HumanMessage(...), AIMessage(...), ...]}
#       然后把这个 "history" 注入到 Prompt 的 MessagesPlaceholder("history") 里。
#
#    b. LLM 生成回复后，LLMChain 会调用 memory.save_context(inputs, outputs)：
#         inputs  = {"question": "..."}     # 也就是你本次的用户输入
#         outputs = {"answer": "..."}       # 也就是模型生成的输出（键名要与 output_key 对应）
#       memory 会把这次问答追加到它的缓冲中：
#         - 以 HumanMessage(content=question) 追加一条“人类消息”
#         - 以 AIMessage(content=answer)     追加一条“AI 消息”
#
# 3) 下一次 predict() 时，memory.load_memory_variables() 又会带上之前累计的历史消息，
#    因而模型可“记住”之前的对话，实现多轮会话记忆。
#
# 4) 默认情况下，这个“记忆”是保存在当前 Python 进程内存里的，重启进程会清空；
#    也不是分用户隔离的（所有调用者共用同一个 memory）。要做“按会话或用户隔离”的话，
#    需根据 session_id / user_id 动态创建/选择不同的 Memory 实例。
# =========================================

# 创建路由
router = APIRouter()

# 1) 初始化 LLM（指向你的本地/内网 Ollama 服务）
llm = OllamaLLM(
    base_url="http://172.16.21.38:11436",
    model="qwen3:0.6b"
)

# 2) 初始化 Memory
# - memory_key="history"：Prompt 中将用这个名字接收历史（和 MessagesPlaceholder 对上）
# - input_key="question"：告诉 Memory“本次用户输入”的键名是什么（用于保存到历史）
# - output_key="answer" ：告诉 Memory“本次模型输出”的键名是什么（用于保存到历史）
# - return_messages=True：以消息列表的形式返回历史，便于 ChatPromptTemplate 正确注入
memory = ConversationBufferMemory(
    memory_key="history",
    input_key="question",
    output_key="answer",
    return_messages=True
)

# 3) 构建“聊天型”Prompt：
# - 使用 MessagesPlaceholder("history") 表示要把历史对话消息直接拼接到这里
# - 最后一条 Human 消息里放入本轮用户输入 {question}
chat_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的中文助手，请根据对话历史进行连贯回答。"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{question}")
])

# 4) 构建对话链
# - output_key="answer"：与上面的 memory.output_key 对齐，方便 memory.save_context
chain = LLMChain(
    llm=llm,
    prompt=chat_prompt,
    memory=memory,
    output_key="answer",
    verbose=False  # 需要时可改为 True 以在控制台打印链路细节
)

@router.get("/chat")
async def chat(question: str = Query(..., description="用户输入的问题")):
    """
    对话接口（带记忆）
    - 请求流程（一次调用）：
      1) 读取历史：memory.load_memory_variables() -> {"history": [消息列表...]}
      2) 注入 Prompt：MessagesPlaceholder("history") 注入历史消息；再注入 {question}
      3) LLM 生成：得到 answer
      4) 写回记忆：memory.save_context({"question": ...}, {"answer": ...})
         -> 把这轮问答追加到内存缓冲里
    """
    try:
        # 触发一次链式调用（会自动读/写记忆）
        answer_text = chain.predict(question=question)

        return JSONResponse(
            content={
                "question": question,
                "answer": answer_text
            }
        )
    except Exception as e:
        return JSONResponse(
            content={
                "error": str(e),
                "trace": traceback.format_exc()
            },
            status_code=500
        )

@router.get("/_memory")
async def peek_memory():
    """
    辅助调试接口：查看当前保存在 ConversationBufferMemory 里的历史消息。
    - 返回两部分：
      1) messages：按顺序列出所有 Human/AI 消息（role + content）
      2) turns   ：统计一共多少轮（1 轮 = 1 次 human + 1 次 ai）
    """
    try:
        # 直接读取底层的消息列表（按你的 LangChain 版本 Field 名称可能为 memory.chat_memory.messages）
        msgs = getattr(memory.chat_memory, "messages", [])

        readable = []
        for m in msgs:
            # 不同版本消息类的 role 字段名称略有差异，这里做个通用兜底
            role = getattr(m, "type", None) or getattr(m, "role", None) or m.__class__.__name__
            content = getattr(m, "content", "")
            readable.append({"role": role, "content": content})

        # 估算“轮数”（人类消息数量或 AI 消息数量的最小值）
        human_cnt = sum(1 for x in msgs if x.__class__.__name__.lower().startswith("human"))
        ai_cnt = sum(1 for x in msgs if x.__class__.__name__.lower().startswith("ai"))
        turns = min(human_cnt, ai_cnt)

        return JSONResponse(
            content={
                "messages": readable,
                "turns": turns
            }
        )
    except Exception as e:
        return JSONResponse(
            content={
                "error": str(e),
                "trace": traceback.format_exc()
            },
            status_code=500
        )

@router.post("/_memory/clear")
async def clear_memory():
    """
    清空当前进程内的对话记忆（所有调用者共享的那份）。
    注意：生产环境通常应该按 user/session 维护多份独立的记忆，而不是全局共用。
    """
    try:
        memory.clear()
        return JSONResponse(content={"ok": True, "message": "memory cleared"})
    except Exception as e:
        return JSONResponse(
            content={
                "error": str(e),
                "trace": traceback.format_exc()
            },
            status_code=500
        )
